#ifndef PATH_PLANNING_ASTAR_H
#define PATH_PLANNING_ASTAR_H

#include <unordered_map>
#include <set>
#include <queue>
#include <string>
#include <algorithm>
#include <iostream>

using namespace std;

/*
    Performs A* search for "State" class that implements the following methods:
        State(const State &state);
        double scoreEstimate(); // Higher is better. Must be greater than or equal to true value.
        vector<State> nextStates();
        string key(); // Must be unique. Used to determine if we've seen this state before.
        string show();
        bool final(); // True if this state is candidate to be the end of a path.
        bool operator<(const State& rhs) const; // Based on score_estimate, needed for priority_queue
 */

// TrajectoryState
template <class State>
class AStar {

private:
    unordered_map<string,State> closedStates;
    unordered_map<string,string> keyToPrevKey;
    unordered_map<string,int> keyToPathLength;
    priority_queue<State> openStates; // priority queue allows lookup of largest element by default
    State bestState;
    double bestStateScore;
    int bestLength;
    bool bestStateFinal;
    bool _finished;
    string initialKey;
public:

    AStar(State initialState) : bestState(initialState) {
        bestStateScore = initialState.scoreEstimate();
        bestLength = 1;
        initialKey = initialState.key(); // 字符串
        bestStateFinal = initialState.final(); // 是否为最后一步，当前node的simulation_step是否为0
        _finished = bestStateFinal; // 搜索是否结束
        // 把起点加入close_set
        closedStates[initialKey] = initialState;
        keyToPathLength[initialKey] = 1;
        // expand node
        vector<State> newStates = initialState.nextStates();
        // 如果没有结束，就是 simulate_steps!=0
        if(!_finished) {
            for (State state : newStates) {
                string k = state.key();
                keyToPrevKey[k] = initialKey;
                keyToPathLength[k] = 2;
                openStates.push(state);
            }
        }
    }


    bool finished() {
        return _finished;
    }


    double score() {
        return bestStateScore;
    }


    vector<State> path() {
        vector<State> result;
        State nextState = bestState; // 搜索结束的node
        string nextKey = bestState.key();
        while(true) {
            result.push_back(nextState);
            if(nextKey == initialKey) { // key等于起点
                reverse(result.begin(), result.end()); // 翻转
                return result;
            }
            nextKey = keyToPrevKey[nextKey]; // unordered_map<string,string> keyToPrevKey;
            nextState = closedStates[nextKey];
        }
    }


    void calculateSeconds(double maxSeconds) {
        if(_finished) {
            return;
        }
        clock_t startClock = clock();
        while(true) {
            calculateSteps(50);
            if(_finished) {
                return;
            }
            clock_t currentClock = clock();
            double elapsedSeconds = double(currentClock - startClock) / CLOCKS_PER_SEC;
            // 如果上一次搜索找不到路径，那么在超时前继续搜索
            if(elapsedSeconds > maxSeconds) {
                return;
            }
        }
    }

    // max_steps = 50
    void calculateSteps(int max_steps) {
        for(int i = 0; i < max_steps; i++) {
            if(_finished) {
                return;
            }
            calculateStep();
        }
    }

    // A_start search
    void calculateStep() {
        if(_finished) {
            return;
        } else if(openStates.empty()) {
            _finished = true;
            return;
        }
        // open_set 不为空
        State state = openStates.top();
        openStates.pop();
        string k = state.key();
        // 如果已经搜索过
        // count函数用以统计key值在unordered_map中出现的次数
        // 实际上，c++ unordered_map不允许有重复的key。因此，如果key存在，则count返回1，如果不存在，则count返回0
        if(closedStates.count(k) > 0) {
            return;
        }
        // 从open_set中删除，然后加入close_set
        closedStates[k] = state;
        double stateScore = state.scoreEstimate();
        
        // bestStateFinal用于判断一个state的simulation_step是不是等于0
        if(bestStateFinal && stateScore < bestStateScore) {
            _finished = true;
            return;
        }

        int statePathLength = keyToPathLength[k];
        bool current_node_final = state.final();
        bool stateBetterThanBest;

        if(current_node_final && !bestStateFinal) { //当前node为搜索的最后一步，上个搜索的最好node不是最后一步
            stateBetterThanBest = true;
        } else if( !current_node_final && bestStateFinal ) { // 当前node为搜索的不是最后一步，上个搜索的最好node是最后一步
            stateBetterThanBest = false;
        } else if(statePathLength > bestLength) { // 当前node的长度大于上个搜索的最好node的长度
            stateBetterThanBest = true;
        } else if( statePathLength < bestLength ) { // 当前node的长度小于上个搜索的最好node的长度
            stateBetterThanBest = false;
        } else {
            stateBetterThanBest = (stateScore > bestStateScore); // 其他一般情况: 当前node的cost大于上个搜索的最好node的cost
        }

        if(stateBetterThanBest) {
            bestState = state;
            bestStateScore = stateScore;
            bestLength = statePathLength;
            bestStateFinal = current_node_final;
        }
        int newLength = statePathLength + 1; // discretization with 1 meter precision along the road
        // expand node
        vector<State> newStates = state.nextStates();
        cout << "New open states for " << k << " :" << endl;
        for (State newState : newStates) {
            string kNew = newState.key();
            if(closedStates.count(kNew) == 0) { // 如果node不在close_set中,就加入
                keyToPrevKey[kNew] = k;
                keyToPathLength[kNew] = newLength;
                openStates.push(newState);
            }
        }
    }
};

#endif
